package com.neucloud.da.edgeside.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.neucloud.da.edgeside.model.constant.EngineCacheConstant;
import com.neucloud.da.edgeside.model.constant.EngineConstant;
import com.neucloud.da.edgeside.model.entity.InferenceEntity;
import com.neucloud.da.edgeside.model.entity.ModelVersionEntity;
import com.neucloud.da.edgeside.model.entity.RunEngineVersion;
import com.neucloud.da.edgeside.model.enums.InferenceTypeEnum;
import com.neucloud.da.edgeside.model.enums.RunEngineEnum;
import com.neucloud.da.edgeside.model.exception.BusinessException;
import com.neucloud.da.edgeside.model.online.InferenceService;
import com.neucloud.da.edgeside.model.online.InferenceServicePredictor;
import com.neucloud.da.edgeside.model.online.InferenceServicePredictorDesc;
import com.neucloud.da.edgeside.model.service.IStorageService;
import io.fabric8.kubernetes.api.model.*;
import io.fabric8.kubernetes.client.KubernetesClient;
import io.fabric8.kubernetes.client.dsl.MixedOperation;
import io.fabric8.kubernetes.client.dsl.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ResourceLoader;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

@Slf4j
@Component
public class InferenceEngineService {
    @Value("${k8s.default.namespace}")
    String namespace;
    @Value("${model.s3.bucket}")
    private String modelBucket;
    @Value("${model.s3.accesskey}")
    private String STORAGE_S3_ACCESSKEY;
    @Value("${model.s3.secretkey}")
    private String STORAGE_S3_SECRETKEY;
    @Value("${model.s3.endpoint}")
    private String STORAGE_S3_ENDPOINT;
    @Value("${image.base.url}")
    private String imageBaseURL;
    @Autowired
    IStorageService storageService;
    @Autowired
    ResourceLoader resourceLoader;
    @Autowired
    KubernetesClient kubernetesClient;
    @Autowired
    RedisTemplate redisTemplate;

    protected boolean createInference(InferenceEntity inferenceEntity, ModelVersionEntity modelVersion) {
        log.info("InferenceService create start inferenceEntity:{} ",inferenceEntity);
        RunEngineVersion runEngineVersion = modelVersion.getRunEngineVersion();
        //获取推理请求响应处理
        String s3IsvcFile;
        if(StrUtil.isNotBlank(inferenceEntity.getTargetProcFile())){
            s3IsvcFile = String.format(EngineConstant.MODE_SCRIPT_PATH, inferenceEntity.getProjectId(),modelVersion.getModelVersionId(),inferenceEntity.getTargetProcFile());
        }else {
            s3IsvcFile = runEngineVersion.getOnlineProcUrl();
        }
        //获取模型前后处理
        String s3ModelHandleFile;
        if(StrUtil.isNotBlank(modelVersion.getTargetProcFile())){
            s3ModelHandleFile = String.format(EngineConstant.MODE_SCRIPT_PATH, modelVersion.getProjectId(),modelVersion.getModelVersionId(),modelVersion.getTargetProcFile());
        }else {
            s3ModelHandleFile = runEngineVersion.getModelProcUrl();
        }
        //获取模型存储路径
        String s3ModelPath = String.format(EngineConstant.MODE_ARTIFACTS_FULL_PATH,modelBucket,inferenceEntity.getProjectId(), modelVersion.getModelVersionId());

        //获取镜像地址
        String imageUrl = imageBaseURL + runEngineVersion.getImageUrl();
        int cpu = inferenceEntity.getCpu();
        int memory = inferenceEntity.getMemory();
        int gpu = inferenceEntity.getGpu();
        //String namespace = inferenceEntity.getProjectId();
        createSecret(namespace);
        createServiceAccount(namespace);
        try {
            MixedOperation<InferenceService, KubernetesResourceList<InferenceService>, Resource<InferenceService>> inferenceServiceClient = kubernetesClient.resources(InferenceService.class);
            InferenceService inferenceService = inferenceServiceClient.load(resourceLoader.getResource("classpath:InferenceService.yaml").getInputStream()).get();
            inferenceService.getMetadata().setName(inferenceEntity.getName());
            inferenceService.getMetadata().setNamespace(namespace);
            Container container = inferenceService.getSpec().getTransformer().getContainers().get(0);
            RunEngineEnum runEngineEnum = RunEngineEnum.getByCode(runEngineVersion.getEngineId());
            //处理Transformer信息
            try {
                handleTransformer(container,runEngineEnum,inferenceEntity,modelVersion,imageUrl,s3ModelHandleFile,s3IsvcFile);
            } catch (Exception e) {
                e.printStackTrace();
                String errorMsg = "创建推理服务失败，请联系平台管理员";
                if(e instanceof BusinessException){
                    errorMsg = e.getMessage();
                }
                log.error("InferenceService handle transformer fail id:{}",inferenceEntity.getInferenceId());
                throw new BusinessException(errorMsg);
            }
            //处理Predictor信息
            handlePredictor(inferenceService.getSpec().getPredictor(),runEngineEnum,s3ModelPath,cpu,memory,gpu);
            //创建推理服务
            inferenceServiceClient.inNamespace(namespace).createOrReplace(inferenceService);

        } catch (Exception e) {
            e.printStackTrace();
            redisTemplate.delete(String.format(EngineCacheConstant.SYN_ONLINE_RES,inferenceEntity.getInferenceId()));
            throw new BusinessException("创建推理服务失败，请联系平台管理员("+inferenceEntity.getInferenceId()+")");
        }
        log.info("InferenceService create success id:{}",inferenceEntity.getInferenceId());
        return true;
    }

    /**
     * @MethodName handleTransformer
     * @Description 处理Transformer信息
     * @Param
     * @param container
     * @param runEngineEnum
     * @param inferenceEntity
     * @param imageUrl
     * @param s3ModelHandleFile
     * @param s3IsvcFile
     * @return boolean
     * @Author luoyueqiang
     * @Date 2022/11/7 14:33
     **/
    private boolean handleTransformer(Container container, RunEngineEnum runEngineEnum, InferenceEntity inferenceEntity, ModelVersionEntity modelVersion, String imageUrl, String s3ModelHandleFile, String s3IsvcFile){
        //替换名称
        container.setName(container.getName().replaceAll("mode-name",inferenceEntity.getName()));
        //替换镜像
        container.setImage(imageUrl);
        //替换环境变量
        for (EnvVar envVar : container.getEnv()) {
            if( envVar.getName().equals("framework")){
                envVar.setValue(runEngineEnum.desc);
            }
            if( envVar.getName().equals("predictor_host")){
                envVar.setValue(envVar.getValue().replaceAll("mode-name",inferenceEntity.getName()).replaceAll("namespace",namespace));
            }
            if( envVar.getName().equals("model_name")){
                envVar.setValue(handleModelName(runEngineEnum,inferenceEntity,modelVersion));
            }
            if( envVar.getName().equals("s3_unified_file")){
                switch (InferenceTypeEnum.getByCode(inferenceEntity.getType())){
                    case ONLINE:
                        envVar.setValue(InferenceTypeEnum.ONLINE.unifyScript);
                        break;
                    case OFFLINE:
                        envVar.setValue(InferenceTypeEnum.OFFLINE.unifyScript);
                        break;
                }
            }
            if( envVar.getName().equals("s3_model_handle_file")){
                envVar.setValue(s3ModelHandleFile);
            }
            if( envVar.getName().equals("s3_isvc_file")){
                envVar.setValue(s3IsvcFile);
            }
            if( envVar.getName().equals("s3_bucket")){
                envVar.setValue(modelBucket);
            }
            if( envVar.getName().equals("s3_access_key")){
                envVar.setValue(STORAGE_S3_ACCESSKEY);
            }
            if( envVar.getName().equals("s3_secret_key")){
                envVar.setValue(STORAGE_S3_SECRETKEY);
            }
            if( envVar.getName().equals("s3_endpoint")){
                envVar.setValue(STORAGE_S3_ENDPOINT);
            }
        }
        //新增环境变量
        if(CollectionUtil.isNotEmpty(inferenceEntity.getEnvParams())){
            for (Map.Entry<String, String> entry : inferenceEntity.getEnvParams().entrySet()) {
                EnvVar envVar = new EnvVar();
                envVar.setName(entry.getKey());
                envVar.setValue(entry.getValue());
                container.getEnv().add(envVar);
            }
        }

        //替换资源
        ResourceRequirements requirements = new ResourceRequirements();
        Map<String, Quantity> limits = new HashMap<>();
        if(inferenceEntity.getGpu() > 0 ){
            limits.put("nvidia.com/gpu",new Quantity(String.valueOf(inferenceEntity.getGpu())));
        }
        if(inferenceEntity.getCpu() > 0){
            limits.put("cpu",new Quantity(String.valueOf(inferenceEntity.getCpu()),"m"));
        }
        if(inferenceEntity.getMemory() > 0){
            limits.put("memory",new Quantity(String.valueOf(inferenceEntity.getMemory()),"M"));
        }
        requirements.setLimits(limits);
        requirements.setRequests(container.getResources().getRequests());
        container.setResources(requirements);
        return true;
    }

    /**
     * @MethodName handlePredictor
     * @Description 处理Predictor信息
     * @Param
     * @param predictor
     * @param runEngineEnum
     * @param s3ModelPath
     * @param cpu
     * @param memory
     * @param gpu
     * @return void
     * @Author luoyueqiang
     * @Date 2022/9/25 11:45
     **/
    private void handlePredictor(InferenceServicePredictor predictor, RunEngineEnum runEngineEnum, String s3ModelPath, int cpu, int memory, int gpu){
        InferenceServicePredictorDesc predictorDesc = predictor.getSklearn();
        predictorDesc.setStorageUri(s3ModelPath);

        ResourceRequirements requirements = new ResourceRequirements();
        Map<String, Quantity> limits = new HashMap<>();
        if( gpu > 0 ){
            limits.put("nvidia.com/gpu",new Quantity(String.valueOf(gpu)));
        }
        if( cpu > 0){
            limits.put("cpu",new Quantity(String.valueOf(cpu),"m"));
        }
        if( memory > 0){
            limits.put("memory",new Quantity(String.valueOf(memory),"M"));
        }
        requirements.setLimits(limits);
        requirements.setRequests(predictorDesc.getResources().getRequests());
        predictorDesc.setResources(requirements);

        switch (runEngineEnum) {
            case TENSORFLOW:
                predictor.setSklearn(null);
                predictor.setTensorflow(predictorDesc);
                break;
            case PYTORCH:
                predictor.setSklearn(null);
                predictor.setPytorch(predictorDesc);
                break;
            case SKLEARN:
                break;
            case XGBOOST:
                predictor.setSklearn(null);
                predictor.setXgboost(predictorDesc);
                break;
        }
    }

    //创建secret
    private void createSecret(String namespace){
        String createStatus = (String) redisTemplate.opsForValue().get(String.format(EngineCacheConstant.SECRET_CACHE,namespace));
        if(StrUtil.isBlank(createStatus)){
            Secret secret = kubernetesClient.secrets().inNamespace(namespace).withName(EngineConstant.MINIO_SECRET).get();
            if(secret == null){
                try {
                    MixedOperation<Secret, KubernetesResourceList<Secret>, Resource<Secret>> resourceClient = kubernetesClient.resources(Secret.class);
                    secret = resourceClient.load(resourceLoader.getResource("classpath:S3Secret.yaml").getInputStream()).get();
                    secret.getMetadata().setNamespace(namespace);
                    String endPoint = STORAGE_S3_ENDPOINT;
                    if(STORAGE_S3_ENDPOINT.indexOf("http://")>=0){
                        endPoint = STORAGE_S3_ENDPOINT.split("http://")[1];
                    }
                    secret.getMetadata().getAnnotations().put("serving.kserve.io/s3-endpoint",endPoint);
                    secret.getStringData().put("AWS_ACCESS_KEY_ID",STORAGE_S3_ACCESSKEY);
                    secret.getStringData().put("AWS_SECRET_ACCESS_KEY",STORAGE_S3_SECRETKEY);
                    secret.getStringData().put("AWS_END_POINT",endPoint);
                    resourceClient.inNamespace(namespace).createOrReplace(secret);
                } catch (IOException e) {
                    e.printStackTrace();
                    throw new BusinessException("创建Secret失败，请联系平台管理员("+namespace+")");
                }
            }
            redisTemplate.opsForValue().set(String.format(EngineCacheConstant.SECRET_CACHE,namespace), "Y");
        }
    }

    //创建SERVICE_ACCOUNT
    private void createServiceAccount(String namespace){
        String createStatus = (String) redisTemplate.opsForValue().get(String.format(EngineCacheConstant.SERVICE_ACCOUNT_CACHE,namespace));
        if(StrUtil.isBlank(createStatus)){
            ServiceAccount serviceAccount = kubernetesClient.serviceAccounts().inNamespace(namespace).withName(EngineConstant.MINIO_SA).get();
            if(serviceAccount == null){
                try {
                    MixedOperation<ServiceAccount, KubernetesResourceList<ServiceAccount>, Resource<ServiceAccount>> resourceClient = kubernetesClient.resources(ServiceAccount.class);
                    serviceAccount = resourceClient.load(resourceLoader.getResource("classpath:S3ServiceAccount.yaml").getInputStream()).get();
                    serviceAccount.getMetadata().setNamespace(namespace);
                    resourceClient.inNamespace(namespace).createOrReplace(serviceAccount);
                } catch (IOException e) {
                    e.printStackTrace();
                    throw new BusinessException("创建ServiceAccount失败，请联系平台管理员("+namespace+")");
                }
            }
            redisTemplate.opsForValue().set(String.format(EngineCacheConstant.SERVICE_ACCOUNT_CACHE,namespace), "Y");
        }
    }

    /**
     * @MethodName handleModelName
     * @Description 处理模型名称
     * @Param
     * @param runEngineEnum
     * @param inferenceEntity
     * @return java.lang.String
     * @Author luoyueqiang
     * @Date 2022/11/7 13:53
     **/
    private String handleModelName(RunEngineEnum runEngineEnum, InferenceEntity inferenceEntity, ModelVersionEntity modelVersion){
        String modelName = inferenceEntity.getName();
        //pytorch特殊处理
        if(runEngineEnum.code.equals(RunEngineEnum.PYTORCH.code)){
            String s3ModelPath = String.format(EngineConstant.MODE_ARTIFACTS_PATH,inferenceEntity.getProjectId(), modelVersion.getModelVersionId());
            try {
                InputStream configInput = storageService.getObject(modelBucket, s3ModelPath + "/config/config.properties");
                Properties properties = new Properties();
                properties.load(configInput);
                Object modelSnapshot = properties.get("model_snapshot");
                if(modelSnapshot == null){
                    throw new BusinessException("模型配置文件(config.properties)未包含model_snapshot字段。");
                }
                JSONObject jsonObj = JSONObject.parseObject((String) modelSnapshot);
                JSONObject models = jsonObj.getJSONObject("models");
                if(models == null){
                    throw new BusinessException("模型配置文件(config.properties)未包含models字段。");
                }
                Set<String> modelNames = models.keySet();
                if(modelNames.size() > 1){
                    throw new BusinessException("模型配置文件值(config.properties)中包含多个模型,系统只支持发布一个模型。");
                }
                for (String name : modelNames) {
                    modelName = name;
                }
            } catch (Exception e) {
                e.printStackTrace();
                throw new BusinessException("未读取到模型配置文件(config.properties),请联系平台管理员");
            }
        }
        return modelName;
    }

}
